from manimlib import *

__curr_dir = os.path.dirname(__file__)
if __curr_dir not in sys.path:
    sys.path.append(__curr_dir)
from MySpring import *


class Spring2Bar(Scene):
    def construct(self):
        axes = ThreeDAxes()
        frame = self.camera.frame
        frame.set_height(12)

        elem_start = [0, 0, 0]
        elem_end = [10, 0, 0]
        spring_elem = Spring3D(elem_start, elem_end, axes)
        cylinder_radius = 0.5
        cylinder = Group(
            Cylinder(radius=cylinder_radius, height=10, color=YELLOW),
            Circle(radius=cylinder_radius).move_to([0, 0, 5]).set_fill(YELLOW, opacity=1).set_style(stroke_width=0),
            Circle(radius=cylinder_radius).shift([0, 0, -5]).set_fill(YELLOW, opacity=1).set_style(stroke_width=0),
        )
        cylinder.rotate(-PI / 2, axis=UP, about_point=ORIGIN).move_to([5, 0, 0]),
        bar_elem = Group(
            cylinder,
            Sphere(radius=spring_elem[1].radius, color=RED).move_to(elem_start),
            Sphere(radius=spring_elem[1].radius, color=RED).move_to(elem_end),
        )

        # 绘制括号标签
        bar_brace = Brace(bar_elem, DOWN)
        bar_brace_text = Tex("L")
        # 更新位置
        always(bar_brace_text.next_to, bar_brace, DOWN)
        bar_brace_group = Group(
            bar_brace,
            bar_brace_text,
        )

        spring_node1 = spring_elem[1]
        spring_node2 = spring_elem[2]
        spring_left_arrow = Arrow(start=spring_node1.get_center() - np.array([1.5, 0, 0]), end=spring_node1)
        spring_left_arrow.set_style(stroke_color=RED)
        always(spring_left_arrow.next_to, spring_node1, LEFT)
        spring_right_arrow = Arrow(start=spring_node2.get_center() + np.array([1.5, 0, 0]), end=spring_node2)
        spring_right_arrow.set_style(stroke_color=RED)
        always(spring_right_arrow.next_to, spring_node2, RIGHT)
        spring_elem_with_arrow = Group(
            spring_elem,
            spring_left_arrow,
            spring_right_arrow,
        )

        bar_node1 = bar_elem[1]
        bar_node2 = bar_elem[2]
        bar_left_arrow = Arrow(start=bar_node1.get_center() - np.array([1.5, 0, 0]), end=bar_node1)
        bar_left_arrow.set_style(stroke_color=RED)
        always(bar_left_arrow.next_to, bar_node1, LEFT)
        bar_right_arrow = Arrow(start=bar_node2.get_center() + np.array([1.5, 0, 0]), end=bar_node2)
        bar_right_arrow.set_style(stroke_color=RED)
        always(bar_right_arrow.next_to, bar_node2, RIGHT)
        bar_elem_with_arrow = Group(
            bar_elem,
            bar_left_arrow,
            bar_right_arrow,
            bar_brace_group,
        )

        spring_to_bar_arrow = Arrow(start=np.array([0, 1, 0]), end=np.array([0, -1, 0]))
        spring_name_text = Text("Spring element")
        bar_name_text = Text("Bar element")
        phase_2_elem_group = Group(
            spring_name_text,
            spring_elem_with_arrow,
            spring_to_bar_arrow,
            bar_name_text,
            bar_elem_with_arrow,
        ).arrange(DOWN)

        phase_2_elem_group.move_to(ORIGIN).shift(RIGHT).shift(RIGHT).shift(RIGHT)
        self.play(ShowCreation(phase_2_elem_group))
        self.wait(2)

        # 创建目标矩阵
        single_spring_k = Tex(
            r"\begin{bmatrix}  F_1  \\  F_2 \end{bmatrix} = k\begin{bmatrix} 1 & -1 \\  -1  &  1  \end{bmatrix} \begin{bmatrix}  \delta{u_1} \\   \delta{u_2}  \end{bmatrix}",
            color=RED,
        )
        single_spring_k.next_to(spring_left_arrow, LEFT)

        self.play(ShowCreation(single_spring_k))
        self.wait(2)

        # ========================================================
        # ========================================================
        # ========================================================

        single_bar_k = Tex(
            r"\begin{bmatrix}  F_1  \\  F_2 \end{bmatrix} = \frac{EA}{L}\begin{bmatrix}1 & -1 \\-1 & 1\end{bmatrix} \begin{bmatrix}  \delta{u_1} \\   \delta{u_2}  \end{bmatrix}",
            color=RED,
        )
        single_bar_k.next_to(bar_left_arrow, LEFT)
        self.play(ShowCreation(single_bar_k))
        self.wait(2)

        # 相机移到杆单元横截面

        self.play(cylinder.animate.rotate(PI / 3, [0, 1, 0]))
        self.wait(2)
        cylinder[1].set_style(fill_color=RED)
        self.wait(3)
        cylinder[1].set_style(fill_color=YELLOW)
        self.play(cylinder.animate.rotate(-PI / 3, [0, 1, 0]))
        self.wait(2)
